# encoding=utf-8
import argparse
import hashlib
import os
import tempfile
import time
import logging
import traceback
import fastapi
import uvicorn
from asyncio import Lock
from fastapi import FastAPI
from starlette.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from logging.handlers import TimedRotatingFileHandler
from DeeplabWrapper import DeeplabGrade
import cv2
from PIL import Image
from io import BytesIO

parser = argparse.ArgumentParser()
parser.add_argument("--path", type=str)
args = parser.parse_args()


if args.path is None:
    logging.error("run path is empty")
    exit(-1)

run_path = args.path

os.makedirs("runtime-cache", exist_ok=True)
os.makedirs("runtime-log", exist_ok=True)
class_colors = {
    1: (255, 0, 0),  # Red
    2: (0, 255, 0),  # Green
    3: (0, 0, 255),  # Blue
    4: (255, 255, 0),  # Yellow
    5: (255, 0, 255)  # Magenta
}


def set_logger():
    logger = logging.getLogger()  # root
    logger.setLevel(logging.INFO)

    file_handler = TimedRotatingFileHandler(
        filename=os.path.join("runtime-log", "app-log"),
        when='D', interval=1
    )
    file_handler.suffix = "%Y-%m-%d.log"

    console_handler = logging.StreamHandler()
    console_handler.setLevel(logging.INFO)

    logging_format = logging.Formatter(
        "%(asctime)s - %(filename)s - %(funcName)s - %(lineno)d - %(levelname)s - %(message)s")

    file_handler.setFormatter(logging_format)
    console_handler.setFormatter(logging_format)

    logger.addHandler(file_handler)
    logger.addHandler(console_handler)


set_logger()
deeplab_grade_model = DeeplabGrade(run_path)


def generate_time_mills():
    return int(round(time.time() * 1000))


def generate_task_id(name: str):
    str_data = name + str(generate_time_mills())
    encoded = hashlib.sha256(str_data.encode()).hexdigest()
    return encoded


app = FastAPI()
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=False,
    allow_methods=["*"],
    allow_headers=["*"]
)

async_lock = Lock()


class DeeplabGradeController:
    @staticmethod
    @app.post("/run/deeplabv3/five")
    async def run(file: fastapi.UploadFile):
        # save to tmp, use PIL to decode
        temp_dir_name = generate_task_id("hello")
        abs_workdir = os.path.join("runtime-cache", temp_dir_name)
        abs_user_path = os.path.join(abs_workdir, "user_input")
        os.makedirs(abs_user_path, exist_ok=True)

        logging.info("save user path: {}".format(abs_user_path))

        file_content = await file.read()
        image = Image.open(BytesIO(file_content))

        temp_file = os.path.basename(tempfile.NamedTemporaryFile(suffix=".png").name)
        abs_file_path = os.path.join(abs_user_path, temp_file)

        image.save(abs_file_path, format='PNG')

        logging.info("write done pred file: {}".format(abs_file_path))

        # run infer
        async with async_lock:
            try:
                start_time = time.perf_counter()
                # get result
                cal = deeplab_grade_model.infer(abs_file_path)

                # save to local for debug
                out_path = os.path.join(abs_workdir, "user_output")
                os.makedirs(out_path, exist_ok=True)
                result_path = os.path.join(out_path, temp_file)

                # save output image
                cv2.imwrite(result_path, cal)

                # print time cost
                end_time = time.perf_counter()
                logging.info("infer cost: {}s".format(end_time - start_time))

                # return
                return fastapi.responses.FileResponse(result_path)

            except Exception as e:
                logging.exception(traceback.format_exc())
                return JSONResponse(
                    status_code=fastapi.status.HTTP_500_INTERNAL_SERVER_ERROR,
                    content={
                        'msg': traceback.format_exc()
                    }
                )


if __name__ == '__main__':
    uvicorn.run(app=app, host="0.0.0.0", port=30012)
